#!/usr/bin/env python
# -*- coding: UTF-8 -*-
'''
@Project ：BigVGAN 
@File ：bigcodec.py
@Author ： z_q_mao
@Date ：24-9-18 下午6:03 
'''

import torch.nn as nn
from vq.codec_encoder import CodecEncoder
from vq.codec_decoder import CodecDecoder

class BigCodecs(nn.Module):
    def __init__(self):
        super().__init__()

        self.encoder = CodecEncoder()
        self.decoder = CodecDecoder()


    def forward(self, x, vqcodec=False):
        '''
        :param x: audio data shape [B, L]
        :return:
            synthesis audio, shape=[B, L]
            vq loss

        '''
        # print("codec in size:", x.size(), x.unsqueeze(1).size())
        vq_emb = self.encoder(x.unsqueeze(1))
        vq_post_emb, vq_code, commit_loss = self.decoder(vq_emb, vq=True)
        if vqcodec:
            return vq_code
        # print(vq_code.size(), wav.size())
        audio_output = self.decoder(vq_post_emb, vq=False) # .detach().cpu().numpy()
        return audio_output, commit_loss, vq_code

    def get_emb(self):
        return self.decoder.quantizer.get_emb()

    def get_codec(self, x):
        vq_emb = self.encoder(x.unsqueeze(1))
        vq_post_emb, vq_code, commit_loss = self.decoder(vq_emb, vq=True)
        return vq_code

    def codec2wav(self, vqcodec):
        # TODO 待调试
        vq_post_emb = self.decoder.vq2emb(vqcodec)
        audio_output = self.decoder(vq_post_emb, vq=False)  # .detach().cpu().numpy()
        return audio_output
